# main.py
# Author: FZU-MCMF <wubo@fzu.edu.cn>
"""
This script processes all YSITE-T.EXP files in the specified folder.
It extracts the temperature and site occupying fractions  (T-SOFs)  data from the TCm-calc.
It then saves the processed data to an Excel file and optionally plots the data.
"""

import os
import re
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def process_exp(file_path: str) -> pd.DataFrame:
    """Extract temperature and site occupying fraction data from YSITE-T.EXP file"""
    data = {}
    current_column = None
    with open(file_path, "r") as file:
        for line in file:
            if "$ PLOTTED COLUMNS ARE" in line:
                match = re.search(r"Y\(FCC,([A-Za-z#0-9]+)\)", line)
                if not match:
                    continue
                current_column = match.group(1)
                if current_column not in data:
                    data[current_column] = {"T": [], "Value": []}
            elif current_column:
                parts = line.split()
                if len(parts) < 2:
                    continue
                try:
                    temperature, value = parts[:2]
                    data[current_column]["T"].append(float(temperature))
                    data[current_column]["Value"].append(np.float32(value))
                except ValueError:
                    continue

    # Combine all extracted data into a single DataFrame
    dataframes = []
    for column_name, values in data.items():
        df = pd.DataFrame(values)
        df.columns = ['T', column_name]
        df = df.drop_duplicates(subset='T')  # Remove duplicate temperature values
        dataframes.append(df.set_index('T'))

    combined_data = pd.concat(dataframes, axis=1).reset_index()
    combined_data = combined_data.iloc[2:-1]  # Remove the first two and last rows
    combined_data.sort_values(by='T', ascending=True, inplace=True)
    combined_data.drop_duplicates(subset='T', inplace=True)
    return combined_data


def plot_data(data: pd.DataFrame, folder_path: str):
    """Plot the processed data and save the plot as a PNG file"""
    markers = ["o", "s", "^", "D", "*", "P"]
    mark_every = 100  # Add markers every 100 units

    plt.figure(figsize=(12, 8))
    for i, column in enumerate(data.columns[1:]):  # Skip the "T" column
        plt.plot(data["T"], data[column], label=column,
                 marker=markers[i % len(markers)], markevery=mark_every)

    plt.legend(loc="center left", bbox_to_anchor=(1, 0.8), fontsize=18)
    plt.xticks(fontsize=20)
    plt.yticks(fontsize=20)
    plt.title("Site Occupying Fraction vs. Temperature", fontsize=20)
    plt.xlabel("Temperature, K", fontsize=20)
    plt.ylabel("SiteOccupying Fraction", fontsize=20)
    plt.tight_layout()
    plt.savefig(os.path.join(folder_path, "YSITE-T.png"))
    plt.close()


def process_exp_files(folder_path: str, generate_plots: bool):
    """Process all YSITE-T.EXP files in the specified folder"""
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            # Check if the file name is "YSITE-T.EXP" (case-insensitive)
            if file.upper() != "YSITE-T.EXP":
                continue
            file_path = os.path.join(root, file)
            # Skip small files
            if os.path.getsize(file_path) < 10 * 1024:
                print(f"Skipping small file: {file_path}")
                continue

            print(f"Processing {file_path}")
            processed_data = process_exp(file_path)
            processed_data.to_excel(os.path.join(root, "YSITE-T.xlsx"), index=False)
            print(f"Data processed and saved")

            if generate_plots:
                plot_data(processed_data, root)
                print("Image saved.")


def get_directory():
    """Get the directory path from the user"""
    while True:
        directory = input("Enter the directory path (leave blank for current directory): ").strip()
        if not directory:
            return os.getcwd()
        elif os.path.isdir(directory):
            return os.path.abspath(directory)
        else:
            print("Invalid directory path. Please try again.")


def main():
    directory = get_directory()
    # generate_plots = input("Do you want to plots? (y/n): ").lower().startswith("y")
    generate_plots = True
    print(f"Scanning {directory}")
    process_exp_files(directory, generate_plots)


VERSION = "0.3"
INFO = f"""
--- YSITE-T.EXP table tool (ver{VERSION}) ---
This program extracts the site occupying fractions (T-SOFs) from the TCm-calc. YSITE-T.EXP,
developed by FZU-MCMF, main developers: Qiao Yang, Gao Min-Liang, Wu Bo*, et al.,
please contact via wubo@fzu.edu.cn, 654489521@qq.com in case.
"""

if __name__ == "__main__":
    print(INFO)
    main()
    print(INFO)
    input(f"Processing complete. Press Enter to exit...")
